#----------------------------------------------------#
#   将单张图片预测、摄像头检测和FPS测试功能
#   整合到了一个py文件中，通过指定mode进行模式的修改。
#----------------------------------------------------#
import time

import cv2
import numpy as np
from PIL import Image

from riversr.model.u_net.unet import Unet

def test():
    unet = Unet()
    video_path = 'http://113.54.7.17:81/rtp/44010200491310000024_44010200491310000024/hls.m3u8'

    capture=cv2.VideoCapture(video_path)

    ref, frame = capture.read()
    if not ref:
        print("读取视频失败，请检查视频路径是否正确！")

    # 读取某一帧
    ref, frame = capture.read()
    if not ref:
        return "读取视频失败，请检查视频路径是否正确！"
    # 格式转变，BGRtoRGB
    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
    # 转变成Image
    frame = Image.fromarray(np.uint8(frame))
    # 进行检测
    frame = np.array(unet.detect_image_pro(frame))
    # RGBtoBGR满足opencv显示格式
    frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
    return frame


def test2():
    unet = Unet()
    video_path = 'http://113.54.7.17:81/rtp/44010200491310000024_44010200491310000024/hls.m3u8'

    capture=cv2.VideoCapture(video_path)

    ref, frame = capture.read()
    if not ref:
        print("读取视频失败，请检查视频路径是否正确！")

    # 读取某一帧
    ref, frame = capture.read()
    if not ref:
        return "读取视频失败，请检查视频路径是否正确！"
    # 格式转变，BGRtoRGB
    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
    # 转变成Image
    frame = Image.fromarray(np.uint8(frame))
    # 进行检测
    level = unet.detect_image_pro_str(frame)

    return level

# if __name__ == '__main__':
#     test()